import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.model_selection import train_test_split
import json
import os
import re

def load_movie_data():
    """加载合并后的电影数据"""
    data_path = os.path.join('vue-dashboard', 'src', 'data', 'movies.json')
    try:
        with open(data_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"加载数据出错: {e}")
        return []

def preprocess_data(movies):
    """预处理电影数据用于模型训练"""
    # 转换为DataFrame
    df = pd.DataFrame(movies)
    
    # 确保必要的列存在
    required_cols = ['directors', 'genres', 'year', 'rating']
    for col in required_cols:
        if col not in df.columns:
            if col == 'year':
                # 尝试从其他字段提取年份
                if 'release_dates' in df.columns:
                    df['year'] = df['release_dates'].apply(
                        lambda x: int(x[0][:4]) if isinstance(x, list) and len(x) > 0 else 2000
                    )
                else:
                    df['year'] = 2000  # 默认值
            else:
                df[col] = None if col == 'rating' else []
    
    # 筛选和清洗数据
    df = df[required_cols].copy()
    
    # 确保年份是数字
    df['year'] = pd.to_numeric(df['year'], errors='coerce').fillna(2000).astype(int)
    
    # 处理评分数据 - 改进版本
    # 首先检查rating是否为字符串类型，如果是，可能包含多个评分或"无评分"文本
    if df['rating'].dtype == 'object':
        # 定义函数来清理评分数据
        def clean_rating(rating):
            if pd.isna(rating):
                return None
            
            # 如果是字符串类型，可能需要额外处理
            if isinstance(rating, str):
                # 检查是否包含"无评分"
                if "无评分" in rating:
                    # 尝试提取第一个数值
                    match = re.search(r'(\d+\.\d+)', rating)
                    if match:
                        return float(match.group(1))
                    return None
                # 尝试直接转换为浮点数
                try:
                    return float(rating)
                except ValueError:
                    return None
            
            # 如果已经是数值类型，直接返回
            return rating
        
        # 应用清理函数
        df['rating'] = df['rating'].apply(clean_rating)
    
    # 转换为数值并填充缺失值
    df['rating'] = pd.to_numeric(df['rating'], errors='coerce')
    
    # 计算有效评分的平均值
    mean_rating = df['rating'].dropna().mean()
    
    # 填充缺失评分
    df['rating'] = df['rating'].fillna(mean_rating)
    
    # 提取主要导演和类型数量
    df['main_director'] = df['directors'].apply(lambda x: x[0] if isinstance(x, list) and len(x) > 0 else 'Unknown')
    df['genre_count'] = df['genres'].apply(lambda x: len(x) if isinstance(x, list) else 0)
    
    # 提取主要类型
    df['main_genre'] = df['genres'].apply(lambda x: x[0] if isinstance(x, list) and len(x) > 0 else 'Unknown')
    
    return df

def train_prediction_model(movies):
    """基于真实数据训练预测模型"""
    # 数据预处理
    df = preprocess_data(movies)
    
    # 填充缺失值
    df = df.fillna({'main_director': 'Unknown', 'main_genre': 'Unknown'})
    
    # 编码分类特征
    director_encoder = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
    genre_encoder = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
    
    director_encoded = director_encoder.fit_transform(df[['main_director']])
    genre_encoded = genre_encoder.fit_transform(df[['main_genre']])
    
    # 年份标准化
    year_scaler = StandardScaler()
    year_scaled = year_scaler.fit_transform(df[['year']])
    
    # 特征组合
    X = np.hstack([
        director_encoded,
        genre_encoded,
        year_scaled,
        df[['genre_count']].values
    ])
    y = df['rating'].values
    
    # 拆分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # 训练模型
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    # 评估模型
    train_score = model.score(X_train, y_train)
    test_score = model.score(X_test, y_test)
    
    print(f"模型训练完成: 训练集R²: {train_score:.4f}, 测试集R²: {test_score:.4f}")
    
    # 计算导演影响力
    directors_impact = calculate_directors_impact(df, director_encoder, model)
    
    # 计算类型影响力
    genres_impact = calculate_genres_impact(df, genre_encoder, model)
    
    return {
        'model': model,
        'director_encoder': director_encoder,
        'genre_encoder': genre_encoder,
        'year_scaler': year_scaler,
        'train_score': train_score,
        'test_score': test_score,
        'directors_impact': directors_impact,
        'genres_impact': genres_impact
    }

def calculate_directors_impact(df, encoder, model):
    """计算各导演对评分的影响力"""
    # 获取所有导演
    directors = df['main_director'].unique()
    impacts = {}
    
    # 计算基准评分
    baseline = df['rating'].mean()
    
    # 获取模型中导演特征的权重
    director_features_count = len(encoder.categories_[0])
    director_coefficients = model.coef_[:director_features_count]
    
    # 计算每个导演的影响力
    for i, director in enumerate(encoder.categories_[0]):
        if i < len(director_coefficients):
            impact = director_coefficients[i]
            if abs(impact) > 0.01:  # 忽略影响微小的导演
                impacts[director] = float(impact)
    
    # 按影响力排序
    return sorted(impacts.items(), key=lambda x: abs(x[1]), reverse=True)[:20]

def calculate_genres_impact(df, encoder, model):
    """计算各类型对评分的影响力"""
    # 计算基准评分
    baseline = df['rating'].mean()
    
    # 获取模型中类型特征的权重
    director_features_count = len(df['main_director'].unique())
    genre_start_idx = director_features_count
    genre_features_count = len(encoder.categories_[0])
    genre_coefficients = model.coef_[genre_start_idx:genre_start_idx+genre_features_count]
    
    # 计算每个类型的影响力
    impacts = {}
    for i, genre in enumerate(encoder.categories_[0]):
        if i < len(genre_coefficients):
            impact = genre_coefficients[i]
            if abs(impact) > 0.01:  # 忽略影响微小的类型
                impacts[genre] = float(impact)
    
    # 按影响力排序
    return sorted(impacts.items(), key=lambda x: abs(x[1]), reverse=True)

def predict_rating(model_data, input_data):
    """使用训练好的模型预测电影评分"""
    # 提取模型和编码器
    model = model_data['model']
    director_encoder = model_data['director_encoder']
    genre_encoder = model_data['genre_encoder']
    year_scaler = model_data['year_scaler']
    
    # 准备预测数据
    director = input_data.get('director', 'Unknown')
    genre = input_data.get('genre', 'Unknown')
    year = input_data.get('year', 2023)
    genre_count = input_data.get('genre_count', 1)
    
    # 编码特征
    director_encoded = director_encoder.transform([[director]])
    genre_encoded = genre_encoder.transform([[genre]])
    year_scaled = year_scaler.transform([[year]])
    
    # 组合特征
    X_pred = np.hstack([
        director_encoded,
        genre_encoded,
        year_scaled,
        [[genre_count]]
    ])
    
    # 预测
    prediction = model.predict(X_pred)[0]
    
    # 确保评分在合理范围内
    prediction = max(min(prediction, 10.0), 1.0)
    
    return prediction

def save_model_data(model_data):
    """保存模型数据到JSON文件供前端使用"""
    # 提取需要保存的数据
    data_to_save = {
        'directors_impact': dict(model_data['directors_impact']),
        'genres_impact': dict(model_data['genres_impact']),
        'model_score': {
            'train': model_data['train_score'],
            'test': model_data['test_score']
        },
        'top_directors': [name for name, _ in model_data['directors_impact'][:10]],
        'top_genres': [name for name, _ in model_data['genres_impact'][:5]]
    }
    
    # 保存到JSON文件
    output_path = os.path.join('vue-dashboard', 'src', 'data', 'model_insights.json')
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(data_to_save, f, ensure_ascii=False, indent=2)
    
    print(f"模型洞察数据已保存到: {output_path}")

def generate_model_insights():
    """生成模型洞察并保存为前端可用的格式"""
    # 加载电影数据
    movies = load_movie_data()
    if not movies:
        print("无法加载电影数据，无法训练模型")
        return None
    
    # 训练模型
    model_data = train_prediction_model(movies)
    
    # 保存模型洞察
    save_model_data(model_data)
    
    return model_data

if __name__ == "__main__":
    generate_model_insights() 